/*
 * Copyright (c) 2009, Luke Emrose
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the
 * distribution.  Neither the name of Luke Emrose Sole Trader nor the
 * names of its contributors may be used to endorse or promote
 * products derived from this software without specific prior written
 * permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/**
 * @file
 * code by Luke Emrose - evolutionarytheory@gmail.com
 *
 * get_micropolygon_id.cpp
 *
 * 3Delight RslPlugin shader to return a unique id as a colour for each micropolyon
 * during a render.
 *
 * Currently returns the unique micropolygon id in red and the unique grid id in green.
 * Future updates may make use of the blue channel of the return colour also.
 */

#include <iostream>
#include <vector>
#include <cassert>

#include <OpenEXR/ImathVec.h>
#include <boost/lexical_cast.hpp>

#include <rx.h>
#include <RslPlugin.h>

// convenience definition
#define rslFunction( x ) RSLEXPORT int x( RslContext* rsl_context, int argc, const RslArg* argv[] )

#define MAX_GRID_SIZE 8192

/**
 * @brief
 * data structure to store an index, and information about previously visited values
 */
struct index_type
{
	int index;

	bool has_value( const Imath::Vec3< float >& val )
	{
		return false;
//		std::vector< Imath::Vec3< float > >::const_iterator i;
//		for( i = values.begin(); i != values.end(); ++i )
//		{
//			if( *i == val )
//			{
//				std::cout << "has value: " << val << std::endl;
//				return true;
//			}
//		}
//
//		std::cout << "doesn't have value: " << val << " -> now being added" << std::endl;
//		values.push_back( val );
//		return false;
	}
//
//private:
//	std::vector< Imath::Vec3< float > > values;
};

extern "C"
{
	/**
	 * @brief
	 * release memory used by the global index values
	 */
	void cleanup_index( RixContext* context, void* data )
	{
		delete static_cast< index_type* >( data );
	}

	/**
	 * @brief
	 * begin the definition of the current grid
	 */
	rslFunction( start_micropoly_id )
	{
		return 0;
	}

	/**
	 * @brief
	 * end the definition of the current grid
	 */
	rslFunction( end_micropoly_id )
	{
		return 0;
	}

	/**
	 * @brief
	 * function to bake a grid of micropolygons
	 */
	rslFunction( get_micropoly_id )
	{
		// try and read limits:gridsize
		RtFloat o_result = 0;
		RxInfoType_t o_result_type = RxInfoFloat;
		RtInt o_result_count = 0;
		RtInt success = RxOption( "limits:gridsize", &o_result, sizeof( RtInt ), &o_result_type, &o_result_count );

		int max_num_grid_verts = MAX_GRID_SIZE;
		if( success == 0 )
		{
			max_num_grid_verts = ( ( int ) o_result ) * 4;
		}

		// get arguments
		RslColorIter ret_arg( argv[ 0 ] );
		
		RslPointIter point_arg( argv[ 1 ] );
		RslNormalIter norm_arg( argv[ 2 ] );
		
		RslColorIter ci_arg( argv[ 3 ] );
		RslColorIter oi_arg( argv[ 4 ] );

		RslFloatIter u_arg( argv[ 5 ] );
		RslFloatIter v_arg( argv[ 6 ] );

		// number of vertices to shade
		int num_vals = argv[ 0 ]->NumValues();

		// allocate storage memory for everything
		std::vector< RtFloat > p( max_num_grid_verts * 3 );
		std::vector< RtFloat > n( max_num_grid_verts * 3 );
		std::vector< RtFloat > ci( max_num_grid_verts * 3 );
		std::vector< RtFloat > oi( max_num_grid_verts * 3 );
		std::vector< RtFloat > u( max_num_grid_verts );
		std::vector< RtFloat > v( max_num_grid_verts );

		// record the first point hit in this grid
		Imath::Vec3< float > first_point;
		first_point.x = ( *point_arg )[ 0 ];
		first_point.y = ( *point_arg )[ 1 ];
		first_point.z = ( *point_arg )[ 2 ];

		// initialize counters
		int point_id = 0;
		int u_size = 0;
		int v_size = 0;

		// storage for a random number per micropolyon
		RtFloat rand_val = ( RtFloat ) rand() / ( RtFloat ) RAND_MAX;
		
		// iterate through all of the vertices of the micropolygons
		for( int iter_id = 0; iter_id < num_vals; ++iter_id )
		{
			int i;

			assert( point_id < num_vals );

			// collect up all of the micropolygon data
			for( int comp = 0; comp < 3; ++comp )
			{
				p[ point_id * 3 + comp ] = ( *point_arg )[ comp ];
				n[ point_id * 3 + comp ] = ( *norm_arg )[ comp ];
				ci[ point_id * 3 + comp ] = ( *ci_arg )[ comp ];
				oi[ point_id * 3 + comp ] = ( *oi_arg )[ comp ];
			}

			++point_id;

			// only store unique u and v values, and then infer the grid size from the number of
			// unique u and v values stored
			RtFloat curr_u = *u_arg;

			for( i = 0; i < u_size; ++i )
			{
				if( u[ i ] == curr_u )
				{
					break;
				}
			}

			if( i == u_size )
			{
				u[ i ] = curr_u;
				assert( u_size < num_vals );
				++u_size;
			}

			RtFloat curr_v = *v_arg;

			for( i = 0; i < v_size; ++i )
			{
				if( v[ i ] == v_size )
				{
					break;
				}
			}

			if( i == v_size )
			{
				v[ i ] = curr_v;
				assert( v_size < num_vals );
				++v_size;
			}

			// move to next sample point
			++point_arg;
			++norm_arg;
			++ci_arg;
			++oi_arg;
			++u_arg;
			++v_arg;
		}

		//============================
		// iterate through each quad
		// numering each one as we go
		//============================

		// current index for the quad to be given an id
		int quad_index = 0;

		// current index for the grid to be given as id
		int grid_index = 0;

		// retrieve the previous quad index from the global memory pool
		RixStorage* global_storage = rsl_context->GetGlobalStorage();

		// thread lock the global storage
		global_storage->Lock();

		// index for the micropolygon
		index_type* quad_index_ptr = static_cast< index_type* >( global_storage->Get( "quad_index" ) );
		if( quad_index_ptr == 0 )
		{
			// create the quad index value for the first time
			quad_index_ptr = new index_type;
			( *quad_index_ptr ).index = 0;
			global_storage->Set( "quad_index", static_cast< void* >( quad_index_ptr ), cleanup_index );
		}
		else
		{
			// only use this value if we have not been here before....
			if( !( *quad_index_ptr ).has_value( first_point ) )
			{
				// ensure the new value is bigger than the current one....
				assert( ( *quad_index_ptr ).index > quad_index );

				// retrieve the previous value of the quad index
				quad_index = ( *quad_index_ptr ).index;
			}
			else
			{
				return 0;
			}
		}

		// index for the grid
		index_type* grid_index_ptr = static_cast< index_type* >( global_storage->Get( "grid_index" ) );
		if( grid_index_ptr == 0 )
		{
			// create the quad index value for the first time
			grid_index_ptr = new index_type;
			( *grid_index_ptr ).index = 0;
			global_storage->Set( "grid_index", static_cast< void* >( grid_index_ptr ), cleanup_index );
		}
		else
		{
			// only use this value if we have not been here before....
			if( !( *quad_index_ptr ).has_value( first_point ) )
			{
				// ensure the new value is bigger than the current one....
				assert( ( *grid_index_ptr ).index > grid_index );

				// retrieve the previous value of the quad index
				grid_index = ( *grid_index_ptr ).index;
			}
			else
			{
				return 0;
			}
		}

		// allocate memory so that it is initially all set to 0 to hold the id's for each quad
		int max_num_quad_ids = ( u_size ) * ( v_size );
		std::vector< RtInt > quad_id( max_num_quad_ids, 0 );

		// iterate through each quad, assigning it a unique index
		for( int curr_v = 0; curr_v < v_size - 1; ++curr_v )
		{
			for( int curr_u = 0; curr_u < u_size - 1; ++curr_u )
			{
				// calculate the id's of the four vertices of each quad
				int vert_id_a = curr_v * u_size + curr_u;
				int vert_id_b = curr_v * u_size + curr_u + 1;
				int vert_id_c = ( curr_v + 1 ) * u_size + curr_u + 1;
				int vert_id_d = ( curr_v + 1 ) * u_size + curr_u;

				quad_id[ vert_id_a ] = quad_index;
				quad_id[ vert_id_b ] = quad_index;
				quad_id[ vert_id_c ] = quad_index;
				quad_id[ vert_id_d ] = quad_index;

				++quad_index;
			}
		}
		++grid_index;

		// set the quad index value ready for the next round
		( *quad_index_ptr ).index = quad_index;

		// set the grid index value ready for the next round
		( *grid_index_ptr ).index = grid_index;

		// unlock the thread lock on the global storage
		global_storage->Unlock();

		// fill in the return values
		for( int iter_id = 0; iter_id < num_vals; ++iter_id )
		{
			( *ret_arg )[ 0 ] = ( RtFloat ) quad_id[ iter_id ];
			( *ret_arg )[ 1 ] = ( RtFloat ) grid_index;
			( *ret_arg )[ 2 ] = 0.0f;
			++ret_arg;
		}

		return 0;
	}

	/**
	 * @brief
	 * define shading functions that can be called from this rsl plugin
	 */
	static RslFunction myFunctions[] =
	{
		{ "void start_micropoly_id(float)", start_micropoly_id, 0, 0 },
		{ "color get_micropoly_id(point,normal,color,color,float,float)", get_micropoly_id, 0, 0 },
		{ "void end_micropoly_id(float)", end_micropoly_id, 0, 0 },
		0
	};

	/**
	 * @brief
	 * register function table with renderer
	 */
	RSLEXPORT RslFunctionTable RslPublicFunctions( myFunctions, 0, 0 );

};  // extern "C"
